
import torch
import numpy as np
from model import FiberVibrationClassifier

def predict_single_sample(model_path, time_seq, space_feat, freq_seq=None):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 加载模型
    model = FiberVibrationClassifier().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()

    # 准备数据
    time_tensor = torch.tensor(time_seq, dtype=torch.float32).unsqueeze(0).to(device)
    space_tensor = torch.tensor(space_feat, dtype=torch.float32).unsqueeze(0).to(device)
    
    if freq_seq is None:
        freq_seq = np.abs(np.fft.fft(time_seq))[:len(time_seq)//2]
    freq_tensor = torch.tensor(freq_seq, dtype=torch.float32).unsqueeze(0).to(device)
    
    # 推理
    with torch.no_grad():
        output, _ = model(time_tensor, space_tensor, freq_tensor)
        pred_class = torch.argmax(output, dim=1).item()
    
    return pred_class

# 示例
if __name__ == "__main__":
    # 随机生成一个测试样本
    time_seq = np.random.randn(1024).astype(np.float32)
    space_feat = np.random.rand(10).astype(np.float32)

    # 模型路径
    model_path = 'fiber_best_model.pth'

    # 预测类别
    pred = predict_single_sample(model_path, time_seq, space_feat)
    print(f"Predicted class: {pred}")
